import dbm
from contextlib import contextmanager, AbstractContextManager
from os import PathLike
from pathlib import Path
from typing import AnyStr, Dict

from storage.kvstore.base import AbstractKVStore
from storage.mutex import Mutex


class DbmKVStore(AbstractKVStore):
    def __init__(self, path: AnyStr | PathLike, *, lock_type=Mutex):
        self._lock_type = lock_type
        self._path = Path(path)
        self._data = None
        self._init()

    def _init(self):
        if self._data is None:
            self._data = dbm.open(str(self._path), 'c')

    def close(self):
        if self._data is not None:
            self._data.close()
            self._data = None

    def __enter__(self):
        if self._data is None:
            self._init()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def get(self, k: bytes, default: bytes = None) -> bytes:
        if k in self._data:
            return self._data[k]
        return default

    def put(self, k: bytes, v: bytes):
        self._data[k] = v

    def update(self, kv_pairs: Dict[bytes, bytes]):
        for k, v in kv_pairs.items():
            self._data[k] = v

    def __delitem__(self, k: bytes):
        del self._data[k]

    def __getitem__(self, item: bytes):
        return self._data[item]

    def __setitem__(self, key: bytes, value: bytes):
        self._data[key] = value

    @contextmanager
    def tx(self) -> AbstractContextManager:
        with self._lock_type(self._path.resolve().with_suffix('.lock')):
            if self._data is None:
                self._init()
                yield self
            else:
                yield self
